"""
测试工具函数
"""
import os
import sys
import numpy as np
import torch

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))


def create_test_game_state():
    """
    创建测试用的游戏状态
    
    Returns:
        游戏状态字典
    """
    return {
        'player_hands': [
            [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 6, 7],  # 玩家0手牌
            [8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 12, 13, 14],  # 玩家1手牌
            [15, 15, 15, 16, 16, 16, 17, 17, 17, 18, 19, 20, 21],  # 玩家2手牌
            [22, 22, 22, 23, 23, 23, 24, 24, 24, 25, 26, 27, 28]  # 玩家3手牌
        ],
        'current_player': 0,
        'last_played_card': 8,
        'played_cards': [8, 15, 22],  # 各玩家打出的牌
        'game_round': 10,
        'tiles_left': 60,
        'player_scores': [1000, 1000, 1000, 1000],
        'player_winds': ['east', 'south', 'west', 'north'],
        'timestamp': 1234567890
    }


def create_test_hand():
    """
    创建测试用的手牌
    
    Returns:
        手牌列表
    """
    return [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 6, 7]


def create_test_features():
    """
    创建测试用的特征向量
    
    Returns:
        特征数组
    """
    # 创建一个简单的特征向量
    features = np.zeros(324, dtype=np.float32)  # 3*108
    
    # 填充一些值
    for i in range(20):
        features[i] = 1.0
    
    return features


def get_test_device():
    """
    获取测试设备
    
    Returns:
        测试设备
    """
    return 'cuda' if torch.cuda.is_available() else 'cpu'


def assert_tensors_equal(tensor1, tensor2, atol=1e-6):
    """
    断言两个张量是否相等
    
    Args:
        tensor1: 第一个张量
        tensor2: 第二个张量
        atol: 绝对容差
    """
    assert torch.allclose(tensor1, tensor2, atol=atol)


def create_test_directory(directory):
    """
    创建测试目录
    
    Args:
        directory: 目录路径
    """
    os.makedirs(directory, exist_ok=True)


def remove_test_directory(directory):
    """
    删除测试目录
    
    Args:
        directory: 目录路径
    """
    if os.path.exists(directory):
        import shutil
        shutil.rmtree(directory)


def create_dummy_model():
    """
    创建一个简单的dummy模型用于测试
    
    Returns:
        简单的模型
    """
    class DummyModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.fc = torch.nn.Linear(324, 6)  # 假设输入是324维，输出是6个动作
        
        def forward(self, x):
            return self.fc(x)
    
    return DummyModel()


def save_dummy_model(model_path):
    """
    保存dummy模型到指定路径
    
    Args:
        model_path: 模型保存路径
    """
    model = create_dummy_model()
    
    # 创建目录
    os.makedirs(os.path.dirname(model_path), exist_ok=True)
    
    # 保存模型
    torch.save(model.state_dict(), model_path)
    
    # 也保存为TorchScript格式
    scripted_model = torch.jit.script(model)
    scripted_model.save(model_path.replace('.pt', '_scripted.pt'))
    
    return model_path
